# linux python 3.10

# pip install insightface -i https://pypi.tuna.tsinghua.edu.cn/simple/
# pip install onnxruntime -i https://pypi.tuna.tsinghua.edu.cn/simple/
# pip install chromadb -i https://pypi.tuna.tsinghua.edu.cn/simple/

# download_path: /root/.insightface/models/buffalo_l
# Downloading /root/.insightface/models/buffalo_l.zip from https://github.com/deepinsight/insightface/releases/download/v0.7/buffalo_l.zip...

import cv2
import insightface
import numpy as np
from insightface.app import FaceAnalysis
import sys 
import time
import chromadb
import os
from typing import TypedDict

class Person(TypedDict):
  Path: str
  Name: str
  ID: str

# Paths to your Indian face images
image_path = os.getcwd()

image_set = [
    Person(Path=f"{image_path}/0.jpg", Name="洋洋", ID="0"),
    Person(Path=f"{image_path}/1.jpg", Name="李宪", ID="1"),
    Person(Path=f"{image_path}/2.jpg", Name="郑凯", ID="2"),
    Person(Path=f"{image_path}/3.jpg", Name="胡哥", ID="3"),
    Person(Path=f"{image_path}/4.jpg", Name="刘浩然", ID="4"),
    Person(Path=f"{image_path}/5.jpg", Name="沙溢", ID="5"),
    Person(Path=f"{image_path}/6.jpg", Name="沈腾", ID="6"),
]

# Initialize face analysis model
app = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])  # Use 'CUDAExecutionProvider' for GPU
app.prepare(ctx_id=-1)  # ctx_id=-1 for CPU, 0 for GPU

# Initialize ChromaDB client and collection
client = chromadb.PersistentClient(path="./face_db")
collection = client.get_or_create_collection(
    name="face_embeddings",
    metadata={"hnsw:space": "cosine"}  # 使用余弦相似度
)
# 余弦相似度是衡量两个向量方向相似性的指标，值域在[-1,1]之间：
# 值为1表示两个向量方向完全相同（完全相似）
# 值为0表示两个向量正交（无相关性）
# 值为-1表示两个向量方向完全相反（完全不相似）

def get_face_embedding(image_path):
    """Extract face embedding from an image"""
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"Could not read image: {image_path}")
    
    faces = app.get(img)
    
    if len(faces) < 1:
        raise ValueError("No faces detected in the image")
    if len(faces) > 1:
        print("Warning: Multiple faces detected. Using first detected face")
    
    return faces[0].embedding

def load_face_embeddings_from_db():
    """Load all face embeddings from the database"""
    for person in image_set:
        try:
            embedding = get_face_embedding(person["Path"])
            collection.add(
                ids=person["ID"],
                embeddings=[embedding],
                metadatas={"name": person["Name"]}
            )
            print(f"Added {person['ID']}: {person['Name']} to database")
        except Exception as e:
            print(f"Failed to add {person['Name']} to database: {str(e)}")

# 余弦相似度
def compare_faces(emb1, emb2, threshold=0.65): 
    """Compare two embeddings using cosine similarity"""
    similarity = np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))
    return similarity, similarity > threshold

# def run():
#     try:
#         emb1 = get_face_embedding("")
#         emb2 = get_face_embedding("")
#         similarity_score, is_same_person = compare_faces(emb1, emb2)
#         print(f"Similarity Score: {similarity_score:.4f}")
#         print(f"Same person? {'YES' if is_same_person else 'NO'}")     
#     except Exception as e:
#         print(f"Error: {str(e)}")

def search_face(input_path, top_k=3, threshold=0.35):
    """
    使用图像路径搜索相似人脸
    
    参数:
        image_path: 要搜索的人脸图像路径
        top_k: 返回的最相似结果数量
        threshold: 相似度阈值，小于此阈值的距离被视为匹配（注意：这里是距离阈值，而非相似度阈值）
    
    返回:
        匹配的人脸列表，每个元素包含ID、姓名、相似度等信息
    """
    # 获取查询图像的人脸嵌入向量
    query_embedding = get_face_embedding(input_path)
    
    # 在Chroma中查询相似向量
    results = collection.query(
        query_embeddings=[query_embedding.tolist()],
        n_results=top_k
    )
    
    # 处理结果
    matches = []
    if not results["ids"]:
        print("未找到相似人脸")
        return matches
    # print(results)
    for i, (face_id, distance, metadata) in enumerate(zip(results["ids"][0], results["distances"][0], results["metadatas"][0])):
        # Chroma返回的是距离，需要转换为相似度
        similarity = 1 - distance
        
        # 检查是否超过阈值（注意：距离小于阈值意味着相似度高于1-阈值）
        # if distance < threshold:
        #     metadata = results["metadatas"][0][i]
        #     matches.append({
        #         "id": face_id,
        #         "name": metadata.get("name", "未知"),
        #         "similarity": similarity,
        #         "distance": distance
        #     })
            # print(f"找到相似人脸 - ID: {face_id}, 姓名: {metadata.get('name', '未知')}, 相似度: {similarity:.4f}")
        # metadata = results["metadata"]
        print(f"找到相似人脸 - ID: {face_id}, 姓名: {metadata['name']}, 相似度: {similarity:.2f}")
        
    return matches


count = collection.count()
if count > 0:
    print(f"loaded existing database, include {count} faces records.")
else:
    print("Creating new database")
    load_face_embeddings_from_db()

if __name__ == "__main__":
    start_time = time.time()
    input_path = f"{image_path}/{sys.argv[1]}"
    search_face(input_path, top_k=3, threshold=0.35)
    end_time = time.time()
    print(f"Total execution time: {end_time - start_time:.4f} seconds")
